import requests
from PIL import Image
from io import BytesIO
import torchvision.transforms as transforms
import torch
from super_gradients.training import models
from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback
import matplotlib.pyplot as plt

# Image can be both uploaded to colab server or by a direct URL
image_path = "https://www.gulfplaceon30a.com/wp-content/uploads/2020/01/yolo1.jpg"

# Load image from url
response = requests.get(image_path)

# Get PIL image
image = Image.open(BytesIO(response.content))

# Prepare preprcoess transformations
# We resize to [640, 640] by COCO's dataset default, which the model was pretrained on.
preprocess = transforms.Compose([
    transforms.Resize([640, 640]),
    transforms.PILToTensor()
])

# Run preprocess on image. unsqueeze for [Batch x Channels x Width x Height] format
transformed_image = preprocess(image).float().unsqueeze(0)

model = models.get("yolox_s", pretrained_weights="coco", num_classes=80)
model.eval()

# Predict using SG model
with torch.no_grad():
    raw_predictions = model(transformed_image)
predictions = YoloPostPredictionCallback(conf=0.1, iou=0.4)(raw_predictions)[0].numpy()


# Visualize results
boxes = predictions[:, 0:4]
plt.figure(figsize=(10, 10))
plt.plot(boxes[:, [0, 2, 2, 0, 0]].T, boxes[:, [1, 1, 3, 3, 1]].T, '.-')
plt.imshow(image.resize([640, 640]))
plt.show()
